#!/usr/bin/env python3
"""
最简单的mask颜色展示脚本
"""

import sys
from PIL import Image
import numpy as np

# 默认文件路径
mask_path = "seg_result_mask.png" if len(sys.argv) < 2 else sys.argv[1]

# 加载mask
mask = Image.open(mask_path)
mask_array = np.array(mask)

# 转换为彩色显示（应用颜色映射）
if len(mask_array.shape) == 2:
    # 基于mask中的实际类别数量自动生成调色板
    import random
    num_classes = int(mask_array.max()) + 1
    palette = []
    # 0类为黑色
    palette.extend([0, 0, 0])
    # 其他类别随机生成颜色
    for cls in range(1, num_classes):
        color = [random.randint(0, 255) for _ in range(3)]
        palette.extend(color)
    # 需要填充到256*3长度
    palette += [0, 0, 0] * (256 - num_classes)
    colored = Image.fromarray(mask_array).convert('P')
    colored.putpalette(palette)
    colored = colored.convert('RGB')
else:
    colored = mask

# 保存并显示
output = "seg_result_colored.png"
colored.save(output)
colored.show()

print(f"✅ 彩色mask已保存: {output}")

